import requests
import pandas as pd
import sqlite3


def get_alpha_unsubmitted(date, region):
    # N小时作为一个取数区间，小数量多批次的获取数据
    period = pd.DateOffset(hours=1)
    tomorrow = date + pd.DateOffset(days=1)
    tomorrow = tomorrow.isoformat()

    start_date = date.isoformat()
    end_date = (date + period).isoformat()

    s = login()

    all_data = []
    while start_date < tomorrow:
        # 步长最大是100， 即便设置成500，返回的 next 也会强制改为100，此举为服务器行为
        step = 100
        offset = 0

        # 按小时计算
        while True:
            url = (
                f"https://api.worldquantbrain.com/users/self/alphas?limit={step}&offset={offset}"
                "&status=UNSUBMITTED%1FIS_FAIL"
                f"&dateCreated>={start_date}-05:00&dateCreated<{end_date}-05:00"
                f"&settings.region={region}"
            )
            print(url)
            response = s.get(url)
            try:
                data = response.json()
            except BaseException:
                # 这里可能需要查看下错误
                # 比如用ipdb
                # import ipdb; ipdb.set_trace()
                s = login()
            try:
                count = data["count"]
            except KeyError as ke:
                print("no data fund")
                break
            if count == 0:
                break
            print(f"{start_date} -> {end_date}, count:", count)
            results = data["results"]
            df = pd.DataFrame(results)
            all_data.append(df)

            next_url = data["next"]
            # 本轮结束，进行下一轮
            if next_url is None:
                break
            # 更新 offset
            offset += step

        start_date = end_date
        end_date = pd.to_datetime(end_date) + period
        end_date = end_date.isoformat()
        print(f"update start_date: {start_date}, end_date: {end_date}===============")

    if len(all_data) > 0:
        df = pd.concat(all_data)
    else:
        df = pd.DataFrame()
    return df


def write_alpha_to_db(dataframe):
    if not dataframe.empty:
        db = sqlite3.connect("./brain_alphas.sqlite3")

        for index, alpha in dataframe.iterrows():
            try:
                selected_alpha = dict(
                    alpha_id=alpha["id"],
                    name=alpha["name"],
                    type=alpha["type"],
                    author=alpha["author"],
                    dateCreated=alpha["dateCreated"],
                    dateSubmitted=alpha["dateSubmitted"],
                    status=alpha["status"],
                    sharpe=alpha["is"]["sharpe"],
                    fitness=alpha["is"]["fitness"],
                    turnover=alpha["is"]["turnover"],
                    margin=alpha["is"]["margin"],
                    longCount=alpha["is"]["longCount"],
                    shortCount=alpha["is"]["shortCount"],
                    decay=alpha["settings"]["decay"],
                    region=alpha["settings"]["region"],
                    universe=alpha["settings"]["universe"],
                    delay=alpha["settings"]["delay"],
                    exp=alpha["regular"]["code"]
                )
                # print(selected_alpha)
                selected_alpha_df = pd.DataFrame(selected_alpha, index=[0])
                selected_alpha_df.to_sql(
                    name="unsubmitted2504", con=db, if_exists="append", index=False
                )
            except TypeError as ke:
                print("some key type error.")



def main():
    # alpha回测开始时间，
    start_date = "2025-04-22"
    # alpha回测结束时间，获取的数据包含这一天的数据，如果要增量获取下一天的数据start_date = "2025-04-03"
    end_date = "2025-04-30"

    REGIONS_LIST = [
        'USA','GLB', 'EUR', 'ASI', 'CHN', 'KOR', 'TWN', 'JPN', 'HKG', 'AMR'
    ]

    days = pd.date_range(start_date, end_date)
    for region in REGIONS_LIST:
        for date in days:
            print(date)
            unsubmitted_alpha_df = get_alpha_unsubmitted(date, region)
            print(f'{date}该日数据开始存入数据库')
            row_count = len(unsubmitted_alpha_df)
            print(f'共计{row_count}行')
            write_alpha_to_db(unsubmitted_alpha_df)
            print(f'{date}数据全部存入完成')


if __name__ == "__main__":
    # 导入自己的平台登录模块
    from machine_lib import login

    main()